package com.example.rtsp2.handler;

import com.example.rtsp2.service.FFmpegService;
import com.example.rtsp2.service.impl.FFmpegServiceImpl;
import jakarta.annotation.PostConstruct;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.BinaryWebSocketHandler;

import javax.annotation.PreDestroy;
import java.io.IOException;
import java.net.SocketTimeoutException;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.*;


/**
 * 处理视频流和WebSocket会话之间的通信
 * @Author: 杜科璇
 * @Description: 视频流处理程序
 * @Date: 2025/9/5 19:07
 * @Version: 1.0
 */

public class VideoStreamHandler extends BinaryWebSocketHandler {
    private static final Logger logger = LoggerFactory.getLogger(VideoStreamHandler.class);

    // 存储所有WebSocket会话
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    // 存储流ID到FFmpeg服务的映射
    private final Map<String, FFmpegService> streamServices = new ConcurrentHashMap<>();

    // 存储会话到流ID的映射
    private final Map<String, String> sessionToStreamId = new ConcurrentHashMap<>();

    // 存储流ID到会话集合的映射
    private final Map<String, CopyOnWriteArraySet<WebSocketSession>> streamSessions = new ConcurrentHashMap<>();

    private final Map<String, Long> sessionLastSendTime = new ConcurrentHashMap<>();
    private final Map<String, Integer> sessionSendErrors = new ConcurrentHashMap<>();
    private static final int MAX_CONSECUTIVE_ERRORS = 5;

    // FFmpeg配置参数
    private final int fps;
    private final int width;
    private final int height;

    // 修改构造函数，使其不需要@Value注解
    public VideoStreamHandler(int fps, int width, int height) {
        this.fps = fps;
        this.width = width;
        this.height = height;
    }
    // 在 VideoStreamHandler 中添加方法，用于检查流是否有活跃会话
    public boolean hasActiveSessions(String streamId) {
        CopyOnWriteArraySet<WebSocketSession> sessionsSet = streamSessions.get(streamId);
        return sessionsSet != null && !sessionsSet.isEmpty();
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        String sessionId = session.getId();
        sessions.put(sessionId, session);
        logger.info("WebSocket连接已建立: {}", sessionId);
    }

    // 修改 afterConnectionClosed 方法，确保正确停止FFmpeg服务
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        String sessionId = session.getId();
        sessions.remove(sessionId);

        // 从流会话映射中移除
        String streamId = sessionToStreamId.remove(sessionId);
        if (streamId != null) {
            CopyOnWriteArraySet<WebSocketSession> streamSessionsSet = streamSessions.get(streamId);
            if (streamSessionsSet != null) {
                streamSessionsSet.remove(session);

                // 如果没有会话监听此流，停止FFmpeg服务
                if (streamSessionsSet.isEmpty()) {
                    logger.info("没有会话订阅流: {}，停止FFmpeg服务", streamId);
                    FFmpegService service = streamServices.get(streamId);
                    if (service != null) {
                        service.stop();
                    }
                }
            }
        }

        logger.info("WebSocket连接已关闭: {}, 状态: {}", sessionId, status);
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) {
        // 处理二进制消息（如果需要）
    }

    // 处理文本消息
    @Override
    protected void handleTextMessage(WebSocketSession session, org.springframework.web.socket.TextMessage message) {
        try {
            String payload = message.getPayload();
            if (payload.startsWith("streamId:")) {
                String streamId = payload.substring("streamId:".length());

                // 检查流是否存在
                if (streamServices.containsKey(streamId)) {
                    // 更新会话到流ID的映射
                    String oldStreamId = sessionToStreamId.put(session.getId(), streamId);

                    // 从旧的流会话中移除
                    if (oldStreamId != null) {
                        CopyOnWriteArraySet<WebSocketSession> oldSessions = streamSessions.get(oldStreamId);
                        if (oldSessions != null) {
                            oldSessions.remove(session);
                        }
                    }

                    // 添加到新的流会话
                    CopyOnWriteArraySet<WebSocketSession> sessionsSet = streamSessions
                            .computeIfAbsent(streamId, k -> new CopyOnWriteArraySet<>());
                    sessionsSet.add(session);

                    // 启动流（如果未运行）
                    FFmpegService service = streamServices.get(streamId);
                    if (!service.isRunning()) {
                        service.start();
                    }

                    logger.info("会话 {} 订阅了流 {}", session.getId(), streamId);
                } else {
                    logger.warn("会话 {} 尝试订阅不存在的流: {}", session.getId(), streamId);
                    session.sendMessage(new org.springframework.web.socket.TextMessage("ERROR:Stream not found"));
                }
            }
        } catch (IOException e) {
            logger.error("处理文本消息时出错", e);
        }
    }

    // 添加新的RTSP流
    public boolean addStream(String streamId, String rtspUrl) {
        if (streamServices.containsKey(streamId)) {
            return false;
        }

        // 使用具体的实现类而不是接口
        FFmpegService service = new FFmpegServiceImpl(rtspUrl, fps, width, height);
        service.setFrameConsumer(frame -> broadcastFrame(streamId, frame));
        streamServices.put(streamId, service);

        logger.info("添加新的RTSP流: {} -> {}", streamId, rtspUrl);
        return true;
    }


    // 移除RTSP流
    // 移除RTSP流
    public boolean removeStream(String streamId) {
        try {
            logger.info("开始移除流: {}", streamId);

            // 获取并停止FFmpeg服务
            FFmpegService service = streamServices.remove(streamId);
            if (service != null) {
                logger.info("停止FFmpeg服务，流ID: {}", streamId);
                service.stop();
            } else {
                logger.warn("FFmpeg服务不存在，流ID: {}", streamId);
            }

            // 移除所有相关会话
            CopyOnWriteArraySet<WebSocketSession> sessionsSet = streamSessions.remove(streamId);

            // 同时清理 sessionToStreamId 中所有指向该流ID的条目
            int removedSessions = 0;
            if (sessionsSet != null && !sessionsSet.isEmpty()) {
                logger.info("移除 {} 个会话，流ID: {}", sessionsSet.size(), streamId);
                for (WebSocketSession session : sessionsSet) {
                    // 从会话到流ID的映射中移除
                    String removedStreamId = sessionToStreamId.remove(session.getId());
                    if (removedStreamId != null) {
                        removedSessions++;
                    }

                    // 发送流移除消息
                    try {
                        if (session.isOpen()) {
                            session.sendMessage(new org.springframework.web.socket.TextMessage("INFO:Stream removed"));
                            session.close(CloseStatus.NORMAL);
                        }
                    } catch (IOException e) {
                        logger.warn("发送流移除消息失败，会话ID: {}", session.getId(), e);
                    }
                }
            } else {
                logger.info("没有找到与会话关联的流，流ID: {}", streamId);

                // 即使sessionsSet为空，也要清理sessionToStreamId中所有指向该流ID的条目
                removedSessions = cleanupSessionToStreamId(streamId);
                logger.info("从sessionToStreamId中移除了 {} 个指向流 {} 的条目", removedSessions, streamId);
            }

            logger.info("成功移除RTSP流: {}, 共清理 {} 个会话关联", streamId, removedSessions);
            return true;

        } catch (Exception e) {
            logger.error("移除流时发生异常，流ID: {}", streamId, e);
            return false;
        }
    }

    // 清理sessionToStreamId中所有指向指定流ID的条目
    private int cleanupSessionToStreamId(String streamId) {
        int removedCount = 0;
        // 使用迭代器安全地移除条目
        Iterator<Map.Entry<String, String>> iterator = sessionToStreamId.entrySet().iterator();
        while (iterator.hasNext()) {
            Map.Entry<String, String> entry = iterator.next();
            if (streamId.equals(entry.getValue())) {
                iterator.remove();
                removedCount++;
            }
        }
        return removedCount;
    }

    // 启动特定流
    public boolean startStream(String streamId) {
        FFmpegService service = streamServices.get(streamId);
        if (service != null && !service.isRunning()) {
            service.start();
            return true;
        }
        return false;
    }

    // 停止特定流
    public boolean stopStream(String streamId) {
        FFmpegService service = streamServices.get(streamId);
        if (service != null && service.isRunning()) {
            service.stop();
            return true;
        }
        return false;
    }

    // 暂停特定流
    public boolean pauseStream(String streamId) {
        FFmpegService service = streamServices.get(streamId);
        if (service != null && service.isRunning()) {
            service.pause();
            return true;
        }
        return false;
    }

    // 恢复特定流
    public boolean resumeStream(String streamId) {
        FFmpegService service = streamServices.get(streamId);
        if (service != null && service.isPaused()) {
            service.resume();
            return true;
        }
        return false;
    }

    // 广播帧到特定流的所有会话
    private void broadcastFrame(String streamId, byte[] frame) {
        CopyOnWriteArraySet<WebSocketSession> sessionsSet = streamSessions.get(streamId);
        if (sessionsSet == null || sessionsSet.isEmpty()) {
            return;
        }

        // 记录发送失败的会话
        Set<WebSocketSession> sessionsToRemove = ConcurrentHashMap.newKeySet();

        // 对每个会话进行智能发送
        for (WebSocketSession session : sessionsSet) {
            String sessionId = session.getId();

            // 检查会话错误计数
            int errorCount = sessionSendErrors.getOrDefault(sessionId, 0);
            if (errorCount >= MAX_CONSECUTIVE_ERRORS) {
                logger.warn("会话 {} 错误过多，暂停发送", sessionId);
                sessionsToRemove.add(session);
                continue;
            }

            // 检查发送频率，避免过快发送
            long lastSendTime = sessionLastSendTime.getOrDefault(sessionId, 0L);
            long currentTime = System.currentTimeMillis();
            long elapsed = currentTime - lastSendTime;

            if (elapsed < (1000 / fps) && elapsed > 0) {
                // 发送太快，跳过此帧
                continue;
            }

            try {
                if (session.isOpen()) {
                    BinaryMessage message = new BinaryMessage(frame);

                    synchronized (session) {
                        session.sendMessage(message);
                    }

                    // 更新最后发送时间
                    sessionLastSendTime.put(sessionId, currentTime);
                    // 重置错误计数
                    sessionSendErrors.put(sessionId, 0);

                } else {
                    sessionsToRemove.add(session);
                }
            } catch (SocketTimeoutException e) {
                logger.warn("发送帧到客户端超时: {}", sessionId);
                sessionSendErrors.put(sessionId, errorCount + 1);
            } catch (IOException e) {
                logger.warn("发送帧到客户端失败: {}", e.getMessage());
                sessionSendErrors.put(sessionId, errorCount + 1);
            } catch (Exception e) {
                logger.error("发送帧时发生未知错误", e);
                sessionSendErrors.put(sessionId, errorCount + 1);
            }
        }

        // 移除所有失败的会话
        if (!sessionsToRemove.isEmpty()) {
            sessionsToRemove.forEach(session -> {
                sessionsSet.remove(session);
                sessionToStreamId.remove(session.getId());
                sessionLastSendTime.remove(session.getId());
                sessionSendErrors.remove(session.getId());

                try {
                    if (session.isOpen()) {
                        session.close(CloseStatus.SERVER_ERROR);
                    }
                } catch (IOException e) {
                    logger.debug("关闭会话时发生错误", e);
                }
            });

            logger.info("移除了 {} 个问题会话，流ID: {}", sessionsToRemove.size(), streamId);
        }
    }

    // 获取活跃会话数量
    public int getActiveSessionCount() {
        return sessions.size();
    }

    // 获取流状态
    public Map<String, Object> getStreamStatus() {
        Map<String, Object> status = new ConcurrentHashMap<>();
        for (Map.Entry<String, FFmpegService> entry : streamServices.entrySet()) {
            Map<String, Object> streamInfo = new ConcurrentHashMap<>();
            FFmpegService service = entry.getValue();

            streamInfo.put("rtspUrl", service.getRtspUrl());
            streamInfo.put("running", service.isRunning());
            streamInfo.put("paused", service.isPaused());

            CopyOnWriteArraySet<WebSocketSession> sessionsSet = streamSessions.get(entry.getKey());
            streamInfo.put("listeners", sessionsSet != null ? sessionsSet.size() : 0);

            status.put(entry.getKey(), streamInfo);
        }
        return status;
    }

    @PreDestroy
    public void cleanup() {
        // 停止所有FFmpeg服务
        for (FFmpegService service : streamServices.values()) {
            service.stop();
        }

        // 关闭所有WebSocket会话
        for (WebSocketSession session : sessions.values()) {
            try {
                session.close();
            } catch (IOException e) {
                logger.warn("关闭WebSocket会话时出错", e);
            }
        }
    }
    /**
     * 健康检查 功能：检测WebSocket会话是否正常
     * */
    @PostConstruct
    public void init() {
        // 启动健康检查任务
        ScheduledExecutorService healthCheckExecutor = Executors.newSingleThreadScheduledExecutor();
        healthCheckExecutor.scheduleAtFixedRate(() -> {
            try {
                checkConnectionHealth();
            } catch (Exception e) {
                logger.error("健康检查异常", e);
            }
        }, 30, 30, TimeUnit.SECONDS); // 每30秒检查一次
    }

    private void checkConnectionHealth() {
        logger.debug("开始连接健康检查");
        int unhealthyCount = 0;

        for (WebSocketSession session : sessions.values()) {
            if (!session.isOpen()) {
                continue;
            }

            String sessionId = session.getId();
            long lastActiveTime = sessionLastSendTime.getOrDefault(sessionId, 0L);
            long inactiveDuration = System.currentTimeMillis() - lastActiveTime;

            // 如果会话超过60秒没有活动，认为不健康
            if (inactiveDuration > 60000) {
                logger.warn("会话 {} 不活跃时间过长: {}ms", sessionId, inactiveDuration);
                unhealthyCount++;

                try {
                    session.close(CloseStatus.SESSION_NOT_RELIABLE);
                } catch (IOException e) {
                    logger.debug("关闭不健康会话时发生错误", e);
                }
            }
        }
        if (unhealthyCount > 0) {
            logger.info("健康检查关闭了 {} 个不健康会话", unhealthyCount);
        }
    }
}